# src/database/connection.py
"""
Postgres connection pool helper using psycopg2.
Provides `init_db_pool()` and context manager `get_conn()`.

Usage:
    from database.connection import init_db_pool, get_conn
    init_db_pool()
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute("SELECT 1")
"""

import os
from contextlib import contextmanager
import logging

import psycopg2
from psycopg2 import pool

logger = logging.getLogger(__name__)

# Read DB config from environment variables with sane defaults
DB_CONFIG = {
    "host": os.getenv("PGHOST", "localhost"),
    "port": int(os.getenv("PGPORT", 5432)),
    "dbname": os.getenv("PGDATABASE", "hms_db"),
    "user": os.getenv("PGUSER", "postgres"),
    "password": os.getenv("PGPASSWORD", "root"),
}

_connection_pool = None


def init_db_pool(minconn: int = 1, maxconn: int = 10):
    """
    Initialize the module-level connection pool.
    Safe to call multiple times.
    """
    global _connection_pool
    if _connection_pool is None:
        try:
            _connection_pool = psycopg2.pool.ThreadedConnectionPool(
                minconn, maxconn, **DB_CONFIG
            )
            logger.info("Initialized Postgres connection pool")
        except Exception as e:
            logger.exception("Failed to initialize DB pool: %s", e)
            raise
    return _connection_pool


@contextmanager
def get_conn():
    """
    Context manager that yields a connection and ensures commit/rollback and
    returns the connection to the pool.
    Example:
        with get_conn() as conn:
            with conn.cursor() as cur:
                cur.execute(...)
    """
    global _connection_pool
    if _connection_pool is None:
        init_db_pool()
    conn = None
    try:
        conn = _connection_pool.getconn()
        yield conn
        conn.commit()
    except Exception:
        if conn:
            conn.rollback()
        raise
    finally:
        if conn:
            _connection_pool.putconn(conn)
